
# 强化学习起飞控制训练脚本（使用 SAC 算法）

import os
import logging

# 引入 Stable-Baselines3 的 SAC 算法及工具组件
from stable_baselines3 import SAC
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback

# 导入起飞增稳环境和 DCS 数据接口
from takeoff_stabilize_env import TakeoffStabilizeEnv
from getDCSdata import get_dcs_data

# ================================
# 日志系统初始化
# ================================

logging.basicConfig(
    level=logging.INFO,  # 日志等级为 INFO
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("TrainTakeoffSAC")

class TakeoffEnvWrapper(TakeoffStabilizeEnv):
    def __init__(self, fast_mode=False):
        super().__init__()
        self.fast_mode = fast_mode

    def step(self, action):
        self._take_action(action)  # 执行动作（控制飞机）

        # 在非快速模式下添加延时，确保与 DCS 交互频率匹配
        if not self.fast_mode:
            import time
            time.sleep(0.1)

        data = get_dcs_data()  # 获取最新传感器数据
        state = self._extract_state(data)  # 提取状态向量

        # 计算奖励和是否达到目标
        reward, done_by_reward = self._calculate_reward(data, action)
        done_by_safety = self._is_done(state, data)  # 根据飞行状态判断是否终止
        done = done_by_reward or done_by_safety  # 只要一个终止条件触发即结束该 step

        self._log_step(reward, data, done)  # 将本步记录写入日志
        self.state = state  # 更新环境状态
        self.step_count += 1  # 步数递增
        return state, reward, done, {}

    def close(self):
        # 释放环境资源（如 DCS 接口），当前为占位
        pass

# ================================
# 主训练函数
# ================================

def main():
    # 设置训练参数
    total_timesteps = 50000  # 总训练步数
    model_name = "sac_takeoff_task3"  # 模型文件名前缀

    # 路径设置
    result_dir = "results"  # 总结果目录
    model_dir = os.path.join(result_dir, "models")  # 模型保存路径
    log_dir = os.path.join(result_dir, "logs")  # 日志目录
    eval_log_dir = os.path.join(log_dir, "eval")  # 评估日志路径
    tensorboard_log = os.path.join(log_dir, "tensorboard")  # TensorBoard 日志路径
    model_path = os.path.join(model_dir, model_name)  # 模型完整保存路径（不含扩展名）

    # 确保所有路径存在
    os.makedirs(model_dir, exist_ok=True)
    os.makedirs(log_dir, exist_ok=True)
    os.makedirs(eval_log_dir, exist_ok=True)

    # 初始化环境（封装为向量环境以适配 SB3）
    logger.info("初始化环境（fast_mode=True）...")
    env = DummyVecEnv([lambda: TakeoffEnvWrapper(fast_mode=True)])

    # 创建 SAC 模型实例
    logger.info("初始化 SAC 模型...")
    model = SAC(
        policy="MlpPolicy",  # 使用多层感知器（全连接神经网络）
        env=env,             # 使用封装好的训练环境
        verbose=1,           # 输出详细训练信息
        tensorboard_log=tensorboard_log,  # 日志保存路径（可用 TensorBoard 查看）
        seed=42              # 设置随机种子，增强实验可重复性
    )

    # 设置训练过程中的模型保存回调
    checkpoint_callback = CheckpointCallback(
        save_freq=10000,         # 每 10000 步保存一次模型
        save_path=model_dir,     # 保存目录
        name_prefix=model_name
    )

    # 设置评估回调函数（定期评估表现并保存最佳模型）
    eval_callback = EvalCallback(
        env,
        best_model_save_path=model_dir,
        log_path=eval_log_dir,
        eval_freq=5000  # 每 5000 步进行一次评估
    )

    # 开始训练过程
    logger.info("开始训练...")
    try:
        model.learn(
            total_timesteps=total_timesteps,  # 指定总训练步数
            callback=[checkpoint_callback, eval_callback]  # 使用回调
        )
    except KeyboardInterrupt:
        # 如果训练被中断，保存中间模型
        logger.warning("训练被中断，保存中间模型...")
        model.save(f"{model_path}_interrupt.zip")
    else:
        # 正常训练结束后保存最终模型
        logger.info("训练完成，保存最终模型...")
        model.save(f"{model_path}.zip")

    logger.info("训练脚本结束。")



if __name__ == "__main__":
    main()
